open Facile
open Easy

open Target
(*open Uav_data*)

let t_ref = ref (Sys.time ())

type variables = {
  alloc: Fd.t array;
  next: Fd.t array;
  costs: Fd.t array;
  cost: Fd.t;
  start_time: Fd.t array;
  end_time: Fd.t array;
  processing_time: Fd.t array;  
}

(* returns variables of problem after posting the constraints *)
let variables = fun goal_type ball video time -> 
   let n = Array.length goal_type and
       m = Array.length ball in

  let alloc = Array.init (n+2*m) (fun _ -> Fd.interval 0 (m-1)) and
      next  = Array.init (n+2*m) (fun _ -> Fd.interval 0 (n+2*m-1)) in

  let start_time      = Array.init (n+2*m) (fun _ -> Fd.interval 0 3600) and
      end_time        = Array.init (n+2*m) (fun _ -> Fd.interval 0 3600) and
      processing_time = Array.init (n+2*m) (fun _ -> Fd.interval 0 3600) in

  (*
  let dur_next = Array.init (n+2*m)
    (fun i ->
      Arith.e2fd (Arith.sum 
      (Array.init m 
	(fun k ->
	  let tki = Array.map (fun t -> Fd.interval t t) time.(k).(i) in
	  fd2e (FdArray.get tki next.(i)) *~ 
	    (fd2e (Reify.boolean (fd2e alloc.(i) =~ i2e k))) 
	)))
    ) in

  (* task duration *)
  for i=0 to n+2*m-1 do
    Cstr.post 
      (fd2e (FdArray.get processing_time next.(i)) =~ fd2e dur_next.(i));
    Cstr.post 
      (fd2e end_time.(i) =~ fd2e start_time.(i) +~ fd2e processing_time.(i));
  done;
  
  (* start times *)
  for k=0 to m-1 do
    Cstr.post (fd2e start_time.(n+k) =~ i2e 0);
  done;
  
  for i=0 to n+m-1 do
    Cstr.post (fd2e end_time.(i) =~ fd2e (FdArray.get start_time next.(i)));
  done;
  *)

  let costs = Array.init m 
    (fun i ->
       let c = Array.init (n+2*m) 
         (fun k -> 
	   fd2e
             (FdArray.get 
		 (Array.map 
		     (fun x -> Fd.interval x x) 
		     time.(i).(k))
		 next.(k)) *~ 
	     (fd2e (Reify.boolean  (fd2e alloc.(k) =~ i2e i )))
         ) in
         Arith.e2fd (Arith.sum c)
    ) in
  
  
  let cost = FdArray.max costs in
  

  
		
  (* instantiate takeoff and landing allocs *)
  for k=0 to m-1 do
    Cstr.post (fd2e alloc.(n+k) =~ i2e k);
    Cstr.post (fd2e alloc.(n+m+k) =~ i2e k);
   done;

  (* each task done only once *)
  Cstr.post (Alldiff.cstr next);

  (* next(k) init constraints *)
  for k=0 to m-1 do
    Cstr.post (fd2e next.(n+m+k) =~  i2e (n+k));
  done;
  for k=0 to (n+2*m-1) do
    Cstr.post (fd2e next.(k) <>~ i2e k);
    for p= 0 to (n+2*m-1) do
      if k<>p then begin
        Cstr.post ((fd2e next.(k) =~ i2e p) =>~~ 
            (fd2e alloc.(k) =~ fd2e alloc.(p)));
      end
    done
  done;

  (* nicolas' vrp constraint *)
  (*Cstr.post (Vrp.cstr n m prev next);*)
  Cstr.post (Nosubtour.cstr n m next);
  
  (* video capacity constraints *)
  for i=0 to m-1 do
    for k=0 to n-1 do
      if (goal_type.(k) = Se || goal_type.(k) = Id) && (video.(i) = 0)
      then 
        Cstr.post (fd2e alloc.(k) <>~ i2e i);
    done
  done;

  (* drop capacity constraints *)
  let drop_allocs_index = (* compute the index of the drop allocs *)
    let index = Array.init n (fun k -> k) in
    Array.of_list 
      (List.filter (fun x-> goal_type.(x) = Dr) 
	 (Array.to_list index))
  in
  let drops_in_plan = Array.init m 
    (fun i ->
       Arith.e2fd
         (Arith.sum_fd
	    (Array.map 
	       (fun tid -> Reify.boolean 
		  (fd2e alloc.(tid) =~ i2e i)) drop_allocs_index)))
  in
  for i=0 to m-1 do
    Cstr.post (fd2e drops_in_plan.(i) <=~ i2e ball.(i))
  done;

  (*
  for k = 0 to m-1 do
    Cstr.post (fd2e costs.(k) =~ fd2e end_time.(n+m+k));
  done;
  *)

  (*Cstr.post (Bound.cstr m alloc costs time);*)

  {
    alloc           = alloc;
    next            = next;
    costs           = costs;
    cost            = cost;
    start_time      = start_time;
    end_time        = end_time;
    processing_time = processing_time; 
  }

(* prints costs of a solution *)
let print_costs = fun variables ->
  let costs = variables.costs and
      st = variables.start_time and
      et = variables.end_time and
      pt = variables.processing_time and
      cost = variables.cost in
  
  let printa = Array.iter (fun c -> Fd.fprint stdout c; Printf.printf "\t%!") in
  (*trace p;*)
  Printf.printf "% 4.2f\t%!" (Sys.time () -. !t_ref); 
  Fd.fprint stdout cost;
  Printf.printf "\t\t\t%!";
  Array.iter (fun c -> Fd.fprint stdout c; Printf.printf "\t%!") costs;
  flush stdout;
  print_newline();
  printa st; print_newline();
  printa pt; print_newline();
  printa et; print_newline();
  print_newline() 

(* save values of an solution *)
let save_goal = fun vars store save ->
  Goals.atomic 
    (fun () ->	  
       let save_in_store = fun () ->
	 store := Some (Fd.elt_value vars.cost,
			Array.map Fd.elt_value vars.alloc,
			Array.map Fd.elt_value vars.next); 
	   save (Array.map Fd.elt_value vars.next)
       in
       match !store with
	   None -> save_in_store ();
	     print_costs vars
	 | Some (c, _, _) ->
	     if Fd.elt_value vars.cost < c 
	     then begin 
	       save_in_store ();
	       print_costs vars
	     end
    )
	    

let alloc_search = fun vars store time ->
  let alloc = vars.alloc in
  let costs = vars.costs in
		  
  Goals.Array.forall
    (Goals.instantiate 
       (Domain.choose 
          (fun a1 a2 -> 
	     (*Invariant.get load.(a1) < Invariant.get load.(a2)*)
	     Fd.max costs.(a1) < Fd.max costs.(a2)
	  ))) alloc

let prev_search = fun vars store time ->
  let alloc = vars.alloc in
  let next = vars.next in
  let choose_index = fun a ->
    let unbounds = Array.map 
      (fun v -> 
	 match Fd.value v with
	     Val _ -> false
	   | _ -> true
      ) a in
    if List.exists (fun x -> x) (Array.to_list unbounds)
    then begin
      let vals = Array.mapi  
	(fun i v ->
	   match Fd.value v with
	       Val _ -> max_int
	     | _ ->  Fd.elt_value alloc.(i)
	) a in
      Utils.arg_min vals
    end
    else raise Not_found in
  Goals.Array.foralli ~select: choose_index
    (fun i -> 
       Goals.instantiate
         (Domain.choose
	    (fun a1 a2 -> 
	       
	       if Fd.is_bound alloc.(i)
	       then
		 let route = Fd.elt_value alloc.(i) in
		 time.(route).(a1).(i) > time.(route).(a2).(i)
	       else Random.bool ()
	    ))) next

let arc_search = fun vars store time ->
  let alloc = vars.alloc in
  let next = vars.next in
  let m = Array.length vars.costs in
  let n = Array.length vars.alloc - (2 * m) in
  let arc = Array.init ((n+2*m)*(n+2*m)) 
    (fun x ->
      let i = x / (n+2*m) and j = x mod (n+2*m) in
      Reify.boolean (fd2e next.(i) =~ i2e j)
    ) in
  (Goals.Array.forall Goals.indomain) arc

(* *)
let local_prev_search = fun vars store time ->
  Goals.Array.forall Goals.indomain vars.next

let heuristical_alloc_search = fun vars store time ->
  let m = Array.length vars.costs in
  let n = Array.length vars.alloc - (2 * m) in
  let insert_cost = fun todo k t ->
    let pos = ref (-1)
    and cost = ref max_int in
    for i=1 to List.length todo.(k) - 2 do
      let ti1 = List.nth todo.(k) (i-1) and
	  ti  = List.nth todo.(k) i in
      let c = time.(k).(ti1).(t) + time.(k).(t).(ti) - time.(k).(ti1).(ti) in
      if c < !cost then begin
	cost := c;
	pos := i;
      end
    done;
    (!pos, !cost)
  in
  let ins = fun todo k pos t ->
    if pos <= 0 then failwith "ins: insert index is wrong!";
    todo.(k) <- Array.to_list 
      (Array.init (List.length todo.(k) +1) 
	  (fun i ->
	    if i<k then List.nth todo.(k) i
	    else if i=k then t
	    else List.nth todo.(k) (i-1)))    
  in
  let insert_k = fun todo k t -> (* insert in agent k task t *)
    let p,_ = insert_cost todo k t in
    ins todo k p t in
  let insert = fun todo t -> (* insert task t with best insertion heuristic *)
     (* allocation *)
    let agent = ref (-1)
    and pos = ref (-1)
    and cost = ref max_int in

    List.iter (fun k ->
    (*for k=0 to m-1 do*)
      let p,c = insert_cost todo k t in
      if c < !cost then begin
	agent := k;
	cost := c;
	pos := p;
      end
    (*done;*)
    ) (Fd.values vars.alloc.(t));
    ins todo !agent !pos t;
    (!agent, !pos, !cost)
  in

  let todo = Array.init m (fun k -> [ n+k; n+m+k ]) in
  
  for i=0 to n-1 do
    if Fd.is_bound vars.alloc.(i) (* Already affected tasks! insert_k *)
    then insert_k todo (Fd.elt_value vars.alloc.(i)) i 
    else begin
      (* allocation *)
      let (k,_,_) = insert todo i in
      Fd.unify vars.alloc.(i) k;
    end
  done


(* complete search goal definition *)
let search_goal = fun vars store time save ->
  (alloc_search vars store time) &&~ (arc_search vars store time)
  &&~ (save_goal vars store save)
  
(* the local optimization search goal *)
let local_search_goal = fun vars store time save ->
  let goal = (alloc_search vars store time) &&~ (prev_search vars store time) &&~ (save_goal vars store save) in
  (((Goals.minimize ~step:2 goal  vars.cost (fun _ -> () )) &&~ 
      Goals.fail) ||~ Goals.success)
    
let complete_search_goal = fun vars store time save ->
  let goal = search_goal vars store time save in
  (((Goals.minimize (goal)  vars.cost (fun _ -> () )) &&~ 
      Goals.fail) ||~ Goals.success)



let select_k_tasks_in_n = fun k distances n ->
  let random_in_list = fun l ->
    let n = List.length l in
    List.nth l (Random.int n) in
  let ft = Random.int n in
  let chosen = ref [ft] in
  let d = 5. (* disrepetency *) in
  let pool = ref (List.filter (fun x -> x<>ft) 
		    (Array.to_list (Array.init n (fun i -> i)))) in
  for i=1 to k-1 do
    let v = random_in_list !chosen in
    let spool = List.sort (fun x y -> compare distances.(v).(x) distances.(v).(y)) !pool in
    let v =
      List.nth spool
      (int_of_float (float (List.length spool) *. ((Random.float 1.) ** d)))
    in
    chosen := v :: !chosen;
    pool := List.filter (fun x-> x<>v) !pool
  done;
  !chosen

let partial_bounding = fun vars store k time ->
  let m = Array.length vars.costs in
  let n = Array.length vars.alloc - (2 * m) in 
  
  match !store with
      None -> failwith "Partial bounding: No store"
    | Some (co, al, pr) ->
	let distances = 
	  Array.init n 
	    (fun i ->
	       Array.init n 
		 (fun j ->
		    time.(0).(i).(j) - time.(0).(j).(j))) in
			  
	let unbounds = 
	  (*Array.to_list (Array.init k (fun i -> Random.int n)) *)
	  select_k_tasks_in_n k distances n
	in
	let prev_unbounds = 
	  let l = ref unbounds in
	  for i = 0 to Array.length pr -1 do
	    if List.mem pr.(i) unbounds then l := i::!l;
	  done;
	  !l in
	
	try
	  for i=0 to Array.length al -1 do
            if not (List.mem i unbounds) 
	    then Fd.unify vars.alloc.(i) al.(i);
            if not (List.mem i prev_unbounds) 
	    then Fd.unify vars.next.(i) pr.(i)
	     (*then Cstr.post (before vars pr.(i) i)*)
          done
        with _ -> ();
	 
	  Cstr.post (fd2e vars.cost <~ i2e (co+1))

let cp_solve = fun 
  ?(save = fun _ -> ()) 
  ?(control = fun _ -> ()) 
  goal_type ball video time trace ->
  let v = variables goal_type ball video time in
  let store = ref None in
  t_ref := Sys.time ();
  ignore(Goals.solve ~control (complete_search_goal v store time save))

let cp2_solove =  fun 
  ?(save = fun _ -> ()) 
  ?(control = fun _ -> ()) 
  goal_type ball video time trace ->
  let v = variables goal_type ball video time in
  let store = ref None in
  t_ref := Sys.time ();
  ignore(Goals.solve ~control (complete_search_goal v store time save))


let stored_cost = fun store ->
  match !store with
    | None -> max_int
    | Some (c,_,_) -> c

let lns_solve = fun goal_type ball video time trace ->
  let v0 = variables goal_type ball video time in
  let store = ref None in
  t_ref := Sys.time ();
  ignore (Goals.solve (search_goal v0 store time (fun _ -> ())));
  while true do

    let vars = variables goal_type ball video time in

    ignore( partial_bounding vars store (4 + Random.int 4) time);
    (*ignore (partial_bounding vars store (2+ Random.int 2) time);*)
    ignore(Goals.solve (complete_search_goal vars store time (fun _ -> ())));
  done

  
